import os
import torch
from torchvision.utils import save_image
from tqdm import tqdm
from picanet_resnet18 import PiCANetResnet18

def test(model, dataloader, device, result_dir="results"):
    model.eval()
    os.makedirs(result_dir, exist_ok=True)

    with torch.no_grad():
        for idx, (images, _) in enumerate(tqdm(dataloader, desc="Testing")):
            images = images.to(device)
            
            # 直接调用模型得到预测（假设模型内部包含了 pred_conv + 插值）
            pred = model(images)  # pred: [B, 1, 224, 224]
            pred = torch.sigmoid(pred)  # 如果模型没加 sigmoid，这里加
            
            for b in range(images.size(0)):
                save_path = os.path.join(result_dir, f"{idx * dataloader.batch_size + b:04d}.png")
                save_image(pred[b], save_path)

